#include <cstdint>

struct WelfordData {
    ::std::uint32_t n;
    float mean;
    float m2;
};

// fk u suiyuan, why you make constexpr global variable not constant (and even contains mangling error)?
#define L1_BUFFER_LEN ::std::uint32_t(128)

#pragma GCC optimize("Ofast")
template<::std::uint32_t thread_num>
__attribute__((global, cooperative))
void kernel_var(
        float const* const __restrict dev_inp,
        float * const __restrict dev_out,
        ::std::uint32_t const nr_elems)
{
    __shared__ ::WelfordData dev_welford_data[thread_num];

    tops_dte_ctx_t ctx;
    // FIXME use dte_scope instead
    ctx.init();
    ::std::uint32_t const thread_idx{blockIdx.x * blockDim.x + threadIdx.x};

    float l1_buf_lhs[L1_BUFFER_LEN];
    tops::mdspan l1_buf_lhs_view(tops::Private, l1_buf_lhs, L1_BUFFER_LEN);

    ::std::uint32_t const elems_per_thread {nr_elems / thread_num};
    ::WelfordData welford_data{};

    ::std::uint32_t const pos_end {thread_idx == thread_num - 1 ? nr_elems : (thread_idx + 1) * elems_per_thread};
    for (::std::uint32_t index{thread_idx * elems_per_thread}; index < pos_end; index += L1_BUFFER_LEN) {
        ::std::uint32_t const chunk {index + L1_BUFFER_LEN < pos_end ? L1_BUFFER_LEN : pos_end - index};
        // Create mdspan for current chunk in L3
        tops::mdspan l3_dev_inp_view(tops::Global, const_cast<float*>(dev_inp + index), chunk);
        // load data from L3 to L1
        tops::memcpy(ctx, l1_buf_lhs_view, l3_dev_inp_view);

        for (::std::uint32_t j{}; j < chunk; ++j) {
            ++welford_data.n;
            float delta {l1_buf_lhs[j] - welford_data.mean};
            welford_data.mean += delta / welford_data.n;
            welford_data.m2 += delta * (l1_buf_lhs[j] - welford_data.mean);
        }
    }

    if constexpr (thread_num == 1) {
        float result {welford_data.m2 / (float(welford_data.n) - 1)};
        ::tops::mdspan result_view(::tops::Private, __builtin_addressof(result), 1);
        ::tops::mdspan dev_out_view(::tops::Global, dev_out, 1);
        ::tops::memcpy(ctx, dev_out_view, result_view);
        return;
    } else {
        ::tops::mdspan l2_dev_welford_single_data_view(::tops::Shared, dev_welford_data + thread_idx, 1);
        ::tops::mdspan l1_dev_welford_single_data_view(::tops::Private, __builtin_addressof(welford_data), 1);
        ::tops::memcpy(ctx, l2_dev_welford_single_data_view, l1_dev_welford_single_data_view);

        __syncblocks();

        if (thread_idx != 0)
        {
            return;
        }

        {
            ::tops::mdspan dev_welford_data_view(::tops::Shared, dev_welford_data, thread_num);
            ::WelfordData l1_dev_welford_data[thread_num];
            ::tops::mdspan l1_dev_welford_data_view(::tops::Private, l1_dev_welford_data, thread_num);
            ::tops::memcpy(ctx, l1_dev_welford_data_view, dev_welford_data_view);

            WelfordData& total = l1_dev_welford_data[0];

#pragma unroll
            for (::std::uint32_t i{1}; i < thread_num; ++i) {
                auto& rhs = l1_dev_welford_data[i];
                float delta = rhs.mean - total.mean;
                ::std::uint32_t n {total.n + rhs.n};
                total.mean = total.mean + delta * rhs.n / n;
                total.m2   = total.m2 + rhs.m2 + delta * delta * total.n * rhs.n / n;
                total.n    = n;
            }

            float result {total.m2 / (float(total.n) - 1)};

            ::tops::mdspan result_view(::tops::Private, __builtin_addressof(result), 1);
            ::tops::mdspan dev_out_view(::tops::Global, dev_out, 1);
            ::tops::memcpy(ctx, dev_out_view, result_view);
        }
    }
}

void GCU_VAR(float * __restrict dev_inp, float * __restrict dev_out, const int nr_elems) {
    if (nr_elems <= 128) {
        ::kernel_var<1><<<1, 1>>>(dev_inp, dev_out, ::std::uint32_t(nr_elems));
    } else if (nr_elems <= 1024) {
        ::kernel_var<8><<<1, 8>>>(dev_inp, dev_out, ::std::uint32_t(nr_elems));
    } else if (nr_elems < 50'000) {
        ::kernel_var<12><<<1, 12>>>(dev_inp, dev_out, ::std::uint32_t(nr_elems));
    } else if (nr_elems < 2'000'000) {
        ::kernel_var<16><<<2, 8>>>(dev_inp, dev_out, ::std::uint32_t(nr_elems));
    } else {
        ::kernel_var<24><<<2, 12>>>(dev_inp, dev_out, ::std::uint32_t(nr_elems));
    }
}
